﻿using System;
using System.Collections.Generic;

using System.IO;
using System.Linq;
using System.Runtime.Serialization.Formatters.Binary;
using System.Text;
using System.Threading;
using System.Xml;

using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;

namespace RayDen.Library.Components.PhotonMapping
{
    public class ModernHashGridPhotonMap : IPhotonMap
    {
        private const float CellSizeScale = 0.001f;
        private Photon[] photons;
        private int[] counts;
        int max_photons;
        internal int photonsCount;
        int prev_scale;
        private float cellSize;
        public AABB bbox;
        public ulong PhotonsCount { get { return (ulong)this.photonsCount; } }


        public Photon[] GetPlainData()
        {
            return this.photons.ToArray();
        }

        public ModernHashGridPhotonMap(int max)
        {
            this.max_photons = max;
            this.prev_scale = 1;
            this.photonsCount = 0;
            this.bbox = new AABB();
            this.photons = new Photon[max];
            this.counts = new int[max];
        }


        public static Vector PhotonDirection(ref Photon p)
        {
            return new Vector(
                DirectionDiscretizationHelper.SinTheta(p.Theta) * DirectionDiscretizationHelper.CosPhi(p.Phi),
                DirectionDiscretizationHelper.SinTheta(p.Theta) * DirectionDiscretizationHelper.SinPhi(p.Phi),
                DirectionDiscretizationHelper.CosTheta(p.Theta)
                );
        }
        public void QueryMap(ref Point pos, float maxDist, int nphotons, out Photon[] result)
        {
            var np = new NearestPhotons();
            np.dist = new float[nphotons + 1];
            np.index = Enumerable.Repeat(new Photon() { Plane = 255 },  1).ToArray();
            np.pos = pos;
            np.max = nphotons;
            np.gotHeap = 0;
            np.found = 0;

            np.dist[0] = maxDist * maxDist;
            this.LocatePhotons(ref np, 1);
            result = np.found > 0 ? np.index.ToArray() : new Photon[0];
        }

        private void LocatePhotons(ref NearestPhotons np, int p)
        {
            Vector hh = (np.pos - bbox.Min) / cellSize;
            int ix = Math.Abs((int)(hh.x));
            int iy = Math.Abs((int)(hh.y));
            int iz = Math.Abs((int)(hh.z));
            var hash = Hash(ix, iy, iz, photonsCount);
            var photon = photons[hash];
            np.found = 1;

  

            if (IsNeighbour(ref np.pos, ref photon.Position, np.dist[0]))
            {
                photon.SetPower(photon.Power*counts[hash]); 
            }

            np.index[0] = photon;
        }

        public Vector IrradianceEstimate(ref Point pos, ref Normal n, float max_dist, int nphotons)
        {
            var normal = n;
            var irrad = new RgbSpectrum(0.0f);
            var np = new NearestPhotons();
            np.dist = new float[nphotons + 1];
            np.index = new Photon[nphotons + 1];
            np.pos = pos;
            np.max = nphotons;
            np.gotHeap = 0;
            np.found = 0;

            np.dist[0] = max_dist * max_dist;

            this.LocatePhotons(ref np, 1);

            if (np.found < 1)
                return new Vector(0);

            for (int index = 0; index < np.index.Length; index++)
            {
                Vector pdir =
                    PhotonDirection(ref np.index[index]);
                var dot = Vector.Dot(ref pdir, ref normal);

                if (dot > 0.0f)
                {
                    irrad = irrad + np.index[index].Power; //*BRDF
                }
            }

            var tmp = (1.0f / MathLab.M_PI) / np.dist[0];

            irrad = irrad * tmp;

            return irrad;

        }

        public void Store(ref RgbSpectrum power, ref Point pos, ref Vector dir)
        {
            if (this.photonsCount >= (this.max_photons - 1))
                return;

            if (power.Length <= 0f)
                return;

            Photon node = this.photons[this.photonsCount];
            node.Power = power;
            node.Position = pos;
            this.bbox = this.bbox.Union(node.Position);

            var dirx = MathLab.Clamp(dir.z, -1f, 1f);
            int theta = Convert.ToInt32(Math.Acos(dirx) * (256.0d / Math.PI));
            if (theta > 255)
                theta = 255;

            node.Theta = (byte)theta;

            int phi = Convert.ToInt32(Math.Atan2(dir.y, dir.x) * (256.0d / 2.0 * Math.PI));


            if (phi > 255)
                phi = 255;
            else

                if (phi < 0)
                    phi = phi + 256;

            node.Phi = (byte)phi;
            this.photons[this.photonsCount] = node;
            this.photonsCount++;

        }

        public void ScalePhotonMap(float scale)
        {
            for (int i = this.prev_scale; i < this.photonsCount; i++)
            {
                this.photons[i].SetPower(this.photons[i].Power * scale);
            }

            this.prev_scale = this.photonsCount;
        }

        public void Balance()
        {

            if (this.photonsCount == 0)
                return;

            var maxRadius2 = 0f;
            var prevPos = photons[0].Position;
            for (int index = 1; index < photons.Length; index++)
            {
                var photon = photons[index];
                var radius = (prevPos - photon.Position).Length;
                if (radius > maxRadius2)
                {
                    maxRadius2 = radius;
                }
                prevPos = photon.Position;
            }

            Photon[] table = new Photon[this.photons.Length];
            cellSize = MathLab.Sqrt(maxRadius2) * CellSizeScale;

            var invCellSize = 1f / cellSize;

            foreach (var point in photons)
            {
                Vector hh = (point.Position - bbox.Min) * invCellSize;
                int ix = Math.Abs((int)(hh.x));
                int iy = Math.Abs((int)(hh.y));
                int iz = Math.Abs((int)(hh.z));
                var hash = Hash(ix, iy, iz, photonsCount);
                table[hash] = point;
                Interlocked.Add(ref counts[hash], 1);
            }

            this.photons = table;
        }

        private static bool IsNeighbour(ref Point p1, ref Point p2, float radius)
        {
            return (p2 - p1).Length <= radius;
        }

        private static uint Hash(int ix, int iy, int iz, int hashGridSize)
        {
            return (uint)(((ix * 73856093U) ^ (iy * 19349663U) ^ (iz * 83492791U)) % hashGridSize);
        }
        private void Swap(ref List<Photon> arr, int a, int b)
        {
            Photon tmp = arr[a];
            arr[a] = arr[b];
            arr[b] = tmp;
        }

        private void MedianSplit(ref List<Photon> p, int start, int end, int median, short axis)
        {
            int left = start;
            int right = end;

            while (right > left)
            {
                double v = p[right].Position[axis];
                int i = left - 1;
                int j = right;
                for (; ; )
                {
                    while (p[++i].Position[axis] < v)
                        ;
                    while (p[--j].Position[axis] > v && j > left)
                        ;
                    if (i >= j)
                        break;
                    this.Swap(ref p, i, j);
                }

                this.Swap(ref p, i, right);
                if (i >= median)
                    right = i - 1;
                if (i <= median)
                    left = i + 1;
            }
        }


        private void BalanceSegment(ref List<Photon> pbal, ref List<Photon> porg, int index, int start, int end)
        {
            int median = 1;
            while ((4 * median) <= (end - start + 1))
                median += median;
            if ((3 * median) <= (end - start + 1))
            {
                median += median;
                median += start - 1;
            }
            else
            {
                median = end - median + 1;
            }

            byte axis = 2;
            Vector dist = this.bbox.Max - this.bbox.Min;
            if ((dist.x > dist.y) && (dist.x > dist.z))
                axis = 0;
            else
                if (dist.y > dist.z)
                    axis = 1;
            this.MedianSplit(ref porg, start, end, median, axis);
            pbal[index] = porg[median];
            pbal[index].SetPlane(axis);

            if (median > start)
            {
                if (start < median - 1)
                {
                    var tmp = this.bbox.Max[axis];
                    this.bbox.Max.SetComponent(axis, pbal[index].Position[axis]);
                    this.BalanceSegment(ref pbal, ref porg, 2 * index, start, median - 1);
                    this.bbox.Max.SetComponent(axis, tmp);
                }
                else
                {
                    pbal[2 * index] = porg[start];
                }
            }

            if (median < end)
            {
                if (end > median + 1)
                {
                    var tmp = this.bbox.Min[axis];
                    this.bbox.Min.SetComponent(axis, pbal[index].Position[axis]);
                    this.BalanceSegment(ref pbal, ref porg, 2 * index + 1, median + 1, end);
                    this.bbox.Min.SetComponent(axis, tmp);
                }
                else
                {
                    pbal[2 * index + 1] = porg[end];
                }
            }
        }

        public void SavePhotonMap(string fileName)
        {/*
            BinaryFormatter frm = new BinaryFormatter();
            using (FileStream fs = new FileStream(fileName, FileMode.Create, FileAccess.ReadWrite))
            {
                BinaryWriter writer = new BinaryWriter(fs);
                string xmlInfo = string.Format(@"<PhotonMapInfo> <photonsCount>{0}</photonsCount> <maxPhotons>{1}</maxPhotons><Half>{2}</Half><PrevScale>{3}</PrevScale>{4}</PhotonMapInfo>",
                                               this.photonsCount, this.max_photons, this.Half, this.prev_scale, this.bbox.ToString());
                byte len = (byte)xmlInfo.Length;
                fs.WriteByte(len);
                fs.Write(Encoding.UTF8.GetBytes(xmlInfo), 0, xmlInfo.Length);
                foreach (Photon photon in this.photons)
                {
                    writer.Write(photon.Position.x);
                    writer.Write(photon.Position.y);
                    writer.Write(photon.Position.z);
                    writer.Write(photon.Power.c1);
                    writer.Write(photon.Power.c2);
                    writer.Write(photon.Power.c3);
                    writer.Write(photon.Phi);
                    writer.Write(photon.Theta);
                    writer.Write(photon.Plane);
                    //                    frm.Serialize(fs, photon);
                }
                fs.Close();
                writer.Dispose();
            }*/
            throw new NotImplementedException();
        }


        public void LoadPhotonMap(string fileName)
        {
            BinaryFormatter frm = new BinaryFormatter();
            using (FileStream fs = new FileStream(fileName, FileMode.Open, FileAccess.Read))
            {
                string xmlInfo = string.Empty;
                int len = fs.ReadByte();
                byte[] xml = new byte[len];
                fs.Read(xml, 0, len);
                xmlInfo = new string(Encoding.UTF8.GetChars(xml));

                XmlDocument param = new XmlDocument();
                param.LoadXml(xmlInfo);
                foreach (XmlNode node in param.ChildNodes[0].ChildNodes)
                {
                    if (node.Name == "photonsCount")
                    {
                        this.photonsCount = int.Parse(node.InnerText.ToString());
                    }

                    if (node.Name == "maxPhotons")
                    {
                        this.max_photons = int.Parse(node.InnerText.ToString());
                    }
                    if (node.Name == "Half")
                    {
                    }
                    if (node.Name == "PrevScale")
                    {
                        this.prev_scale = int.Parse(node.InnerText.ToString());
                    }
                    if (node.Name == "BBox")
                    {
                        string nodeVal = node.InnerText.ToString();
                    }
                }

                this.photons = new Photon[max_photons];
                var sp = fs.Position;
                using (BinaryReader reader = new BinaryReader(fs))
                    for (int index = 0; index < this.photons.Length; index++)
                    {
                        var position = new Point(reader.ReadSingle(), reader.ReadSingle(), reader.ReadSingle());
                        var power = new RgbSpectrum(reader.ReadSingle(), reader.ReadSingle(), reader.ReadSingle());
                        byte phi = reader.ReadByte();
                        byte theta = reader.ReadByte();
                        byte plane = reader.ReadByte();
                        this.photons[index] = new Photon() { Position = position, Power = power, Phi = phi, Theta = theta, Plane = plane };
                    }

                fs.Close();
            }
        }
    }
}